import pandas as pd
import matplotlib.pyplot as plt
import datetime as dt
import statsmodels.api as sm

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
data = pd.read_csv("data.csv", index_col="Date")
data.index = [dt.datetime.strptime(x, "%Y-%m-%d") for x in data.index]
# print(data.head())
# data.plot()
# plt.ylabel("a")
# plt.show()
x = data["沪深300"].values
X = sm.add_constant(x)
y = data["中国平安"].values
model = sm.OLS(y, X)
result = model.fit()
# print(result.params)
plt.figure(figsize=(10, 6))
plt.plot(x, y, 'o', label='中国平安-沪深300')
# 普通最小二乘法
plt.plot(x, result.fittedvalues, 'r--', label='ordinary least square')
plt.legend()
plt.xlabel('沪深300')
plt.ylabel('中国平安')
plt.grid(True)
plt.show()
